import oneflow as torch
import oneflow.nn as nn


class TimeDropout(nn.Module):
    def __init__(self, p):
        super().__init__()

        self.p = p

    def forward(self, inputs):
        """Dropout along Time sequence
        Args:
            inputs: [b, t, v]
        """
        if not self.training:
            return inputs
        else:
            b, t, _ = inputs.size()
            mask = ~(torch.rand(b, t).to(inputs.device) < self.p)
            return inputs * mask.unsqueeze(-1).float()


class FreqDropout(nn.Module):
    def __init__(self, p):
        super().__init__()

        self.p = p

    def forward(self, inputs):
        """Dropout along Time sequence
        Args:
            inputs: [b, t, v]
        """
        if not self.training:
            return inputs
        else:
            b, _, v = inputs.size()
            mask = ~(torch.rand(b, v).to(inputs.device) < self.p)
            return inputs * mask.unsqueeze(-1).float()


class InsideTwoDimDropout(nn.Module):
    def __init__(self, tp, fp):
        super(InsideTwoDimDropout, self).__init__()

        self.tp = tp
        self.fp = fp

    def forward(self, inputs):
        """
        Args:
            x: [batch_size, time_step, feature_dim]
        """
        if not self.training:
            return inputs

        b, t, f = inputs.size()

        time_mask = ~(torch.rand(b, t).to(inputs.device) < self.tp)
        inputs = inputs * time_mask.unsqueeze(-1).float()

        freq_mask = ~(torch.rand(b, f).to(inputs.device) < self.fp)
        inputs = inputs * freq_mask.unsqueeze(-1).float()

        return inputs   


class HeadDropout(nn.Module):
    def __init__(self, p=0.1):
        super(HeadDropout, self).__init__()

        self.p = 0.1

    def forward(self, inputs):
        """
        Args:
            x: [batch_size, n_heads, time1, time2]
        """
        if not self.training:
            return inputs

        b, n, _, _ = inputs.size()

        head_mask = ~(torch.rand(b, n).to(inputs.deice) < self.p)
        inputs = inputs * head_mask.reshape(b, n, _, _).float()

        return inputs
